/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <iostream>
#include <cstring>
#include <random>
#include <acl/acl.h>
#include "securec.h"
#include "atb/atb_infer.h"
#include "chatglm6b_add_operation.h"

static aclError CheckAcl(aclError ret)
{
    if (ret != ACL_ERROR_NONE) {
        std::cerr << __FILE__ << ":" << __LINE__ << " aclError:" << ret << std::endl;
    }
    return ret;
}

static void GenerateTestData(uint8_t *xHost, uint8_t *yHost, uint8_t *goldenHost)
{
    const int totalLengthImm = 16384;
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> dis(1, 50); // 生成1-50之间的随机数
    for (int i = 0; i < totalLengthImm; ++i) {
        uint16_t x = static_cast<uint16_t>(dis(gen));
        uint16_t y = static_cast<uint16_t>(dis(gen));
        uint16_t golden = x + y;
        error_t err = memcpy_s(xHost + i * sizeof(uint16_t), sizeof(uint16_t), &x, 1);
        if (err != 0) {
            std::cout << "[ERROR]: xHost memcpy_s fail!" << std::endl;
        }
        err = memcpy_s(yHost + i * sizeof(uint16_t), sizeof(uint16_t), &y, 1);
        if (err != 0) {
            std::cout << "[ERROR]: yHost memcpy_s fail!" << std::endl;
        }
        err = memcpy_s(goldenHost + i * sizeof(uint16_t), sizeof(uint16_t), &golden, 1);
        if (err != 0) {
            std::cout << "[ERROR]: goldenHost memcpy_s fail!" << std::endl;
        }
    }
}

static void SetCurrentDevice()
{
    const int deviceId = 0;
    std::cout << "[INFO]: aclrtSetDevice " << deviceId << std::endl;
    int ret = aclrtSetDevice(deviceId);
    if (ret != 0) {
        std::cout << "[ERROR]: aclrtSetDevice fail, error:" << ret << std::endl;
        return;
    }
    std::cout << "[INFO]: aclrtSetDevice success" << std::endl;
}

static void CreateTensor(atb::Tensor &tensor)
{
    tensor.dataSize = 16384 * sizeof(uint16_t); // 16384个测试数据
    int ret = aclrtMalloc(&tensor.deviceData, tensor.dataSize, ACL_MEM_MALLOC_HUGE_FIRST);
    if (ret != 0) {
        std::cout << "[ERROR]: aclrtMalloc fail, ret:" << ret << std::endl;
    }
}

static void CreateVariantPack(atb::VariantPack &variantPack, uint8_t *xHost, uint8_t *yHost)
{
    const size_t blockDim = 8;
    size_t inputByteSize = blockDim * 2048 * sizeof(uint16_t);  // uint16_t represent half
    
    atb::TensorDesc desc;
    desc.dtype = ACL_FLOAT;
    desc.format = ACL_FORMAT_NCHW;
    desc.shape.dimNum = 1;
    desc.shape.dims[0] = 16384; // 16384个测试数据

    const int inTensorCount = 2;
    variantPack.inTensors.resize(inTensorCount);

    atb::Tensor &inTensor0 = variantPack.inTensors.at(0);
    inTensor0.desc = desc;
    CreateTensor(inTensor0);
    CheckAcl(aclrtMemcpy(inTensor0.deviceData, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    atb::Tensor &inTensor1 = variantPack.inTensors.at(1);
    inTensor1.desc = desc;
    CreateTensor(inTensor1);
    CheckAcl(aclrtMemcpy(inTensor1.deviceData, inputByteSize, yHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    variantPack.outTensors.resize(1);
    atb::Tensor &outTensor0 = variantPack.outTensors.at(0);
    outTensor0.desc = desc;
    CreateTensor(outTensor0);
}

static void FreeTensor(atb::Tensor &tensor)
{
    if (tensor.deviceData) {
        int ret = aclrtFree(tensor.deviceData);
        if (ret != 0) {
            std::cout << "[ERROR]: aclrtFree fail" << std::endl;
        }
        tensor.deviceData = nullptr;
        tensor.dataSize = 0;
    }
}

static void FreeTensors(atb::SVector<atb::Tensor> &inTensors, atb::SVector<atb::Tensor> &outTensors)
{
    for (size_t i = 0; i < inTensors.size(); ++i) {
        FreeTensor(inTensors.at(i));
    }
    for (size_t i = 0; i < outTensors.size(); ++i) {
        FreeTensor(outTensors.at(i));
    }
}

static bool CompareResult(uint8_t *goldenHost, uint8_t *zHost, size_t dataSize)
{
    if (std::memcmp(goldenHost, zHost, dataSize) == 0) {
        return true; // 如果比较成功，返回true
    } else {
        return false; // 如果比较失败，返回false
    }
}

int main(int argc, const char *argv[])
{
    const uint32_t blockDim = 8;
    const size_t tilingSize = 2 * sizeof(uint32_t);
    const uint32_t one_repeat_calcount = 128;  // fixed
    const uint32_t block_dim_imm = 8;
    const uint32_t tile_num_imm = 8;
    const uint32_t double_buffer_imm = 2;  // fixed
    const uint32_t total_length_imm = block_dim_imm * one_repeat_calcount * tile_num_imm * double_buffer_imm;
    const size_t inputByteSize = blockDim * 2048 * sizeof(uint16_t);
    const size_t outputByteSize = blockDim * 2048 * sizeof(uint16_t);

    uint8_t *xHost = nullptr;
    uint8_t *yHost = nullptr;
    uint8_t *goldenHost = nullptr;

    SetCurrentDevice();

    CheckAcl(aclrtMallocHost(reinterpret_cast<void**>(&xHost), inputByteSize));
    CheckAcl(aclrtMallocHost(reinterpret_cast<void**>(&yHost), inputByteSize));
    CheckAcl(aclrtMallocHost(reinterpret_cast<void**>(&goldenHost), outputByteSize));
    std::cout << "[INFO]: complete aclrtMallocHost!" << std::endl;
    
    GenerateTestData(xHost, yHost, goldenHost);
    std::cout << "[INFO]: complete generate_test_data!" << std::endl;

    atb_speed::ChatGlm6BAddOperation *addOp = new atb_speed::ChatGlm6BAddOperation("addOperation");
    std::cout << "[INFO]: complete CreateOp!" << std::endl;

    atb::VariantPack variantPack;
    CreateVariantPack(variantPack, xHost, yHost);
    std::cout << "[INFO]: complete CreateVariantPack!" << std::endl;

    uint64_t workspaceSize = 0;
    atb::Context *context = nullptr;
    int ret = atb::CreateContext(&context);
    void *stream = nullptr;
    ret = aclrtCreateStream(&stream);
    if (ret != 0) {
        std::cout << "[ERROR]: AsdRtStreamCreate fail, ret:" << ret << std::endl;
        return 1;
    }
    context->SetExecuteStream(stream);
    std::cout << "[INFO]: complete CreateContext!" << std::endl;
    atb::Status st = addOp->Setup(variantPack, workspaceSize, context);
    if (st != 0) {
        std::cout << "[ERROR]: AddOperation setup fail" << std::endl;
        return -1;
    }
    void *workspace = nullptr;
    if (workspaceSize > 0) {
        aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    }
    std::cout << "[INFO]: AddOperation setup success" << std::endl;

    std::cout << "[INFO]: AddOperation execute start" << std::endl;
    st = addOp->Execute(variantPack, (uint8_t*)workspace, workspaceSize, context);
    if (st != 0) {
        std::cout << "[ERROR]: AddOperation execute fail" << std::endl;
        return -1;
    }
    std::cout << "[INFO]: AddOperation setup success" << std::endl;

    // compare results
    uint8_t *zHost = nullptr;
    CheckAcl(aclrtMallocHost((void**)(&zHost), outputByteSize));
    CheckAcl(aclrtMemcpy(zHost, outputByteSize, variantPack.outTensors.at(0).deviceData,
        outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    if (CompareResult(goldenHost, zHost, outputByteSize)) {
        std::cout << "EQUAL: outputs are identical!" << std::endl;
    } else {
        std::cout << "UNEQUAL: outputs are not the same!" << std::endl;
    }

    FreeTensors(variantPack.inTensors, variantPack.outTensors);
    st = atb::DestroyContext(context);
    CheckAcl(aclrtFreeHost(xHost));
    CheckAcl(aclrtFreeHost(yHost));
    CheckAcl(aclrtFreeHost(goldenHost));
    CheckAcl(aclrtFreeHost(zHost));
    CheckAcl(aclrtDestroyStream(stream));
    CheckAcl(aclrtResetDevice(0));
    CheckAcl(aclFinalize());
    return atb::ErrorType::NO_ERROR;
}
